//! Symbol table for NRC language

use crate::ast::*;
use crate::error::{CompilerError, Result};
use std::collections::HashMap;

use super::{
    ConstantSymbol, EnumSymbol, FunctionSymbol, ImportSymbol, StructSymbol, Symbol, TraitSymbol,
    TypeSymbol, VariableSymbol,
};

/// Symbol table for managing symbols in different scopes
pub struct SymbolTable {
    /// Global symbols keyed by name (may contain multiple overloads)
    global_symbols: HashMap<String, Vec<Symbol>>,
    /// Current scope symbols
    current_scope: Vec<HashMap<String, Vec<Symbol>>>,
    /// Current scope level
    current_level: usize,
}

impl SymbolTable {
    /// Create a new symbol table
    pub fn new() -> Self {
        let mut symbol_table = Self {
            global_symbols: HashMap::new(),
            current_scope: vec![HashMap::new()],
            current_level: 0,
        };

        // Initialize built-in functions
        symbol_table.initialize_builtin_functions();
        symbol_table
    }

    /// Initialize built-in functions
    fn initialize_builtin_functions(&mut self) {
        // println function - takes variable arguments and returns void
        let println_func = FunctionDecl {
            name: "println".to_string(),
            generics: vec![],
            parameters: vec![], // Variable arguments - we'll handle this specially
            return_type: None,  // void return type
            body: None,         // External function
            attributes: vec![],
            location: crate::error::Location::new(0, 0, 0),
        };

        let mut println_signature = FunctionSignature::from_decl(&println_func);
        println_signature.is_variadic = true;
        let println_symbol = Symbol::Function(FunctionSymbol {
            name: "println".to_string(),
            declaration: println_func,
            signature: println_signature,
            scope_level: 0,
            exported: true,
        });

        let _ = self.insert_global_symbol("println".to_string(), println_symbol);

        // print function - similar to println but without newline
        let print_func = FunctionDecl {
            name: "print".to_string(),
            generics: vec![],
            parameters: vec![],
            return_type: None,
            body: None,
            attributes: vec![],
            location: crate::error::Location::new(0, 0, 0),
        };

        let mut print_signature = FunctionSignature::from_decl(&print_func);
        print_signature.is_variadic = true;
        let print_symbol = Symbol::Function(FunctionSymbol {
            name: "print".to_string(),
            declaration: print_func,
            signature: print_signature,
            scope_level: 0,
            exported: true,
        });

        let _ = self.insert_global_symbol("print".to_string(), print_symbol);

        // len function - returns length of string/array
        let len_func = FunctionDecl {
            name: "len".to_string(),
            generics: vec![],
            parameters: vec![crate::ast::types::Parameter {
                name: "value".to_string(),
                param_type: Type::Basic(crate::ast::types::BasicType::String),
                default_value: None,
                location: crate::error::Location::new(0, 0, 0),
            }],
            return_type: Some(Type::Basic(crate::ast::types::BasicType::Int)),
            body: None,
            attributes: vec![],
            location: crate::error::Location::new(0, 0, 0),
        };

        let len_signature = FunctionSignature::from_decl(&len_func);
        let len_symbol = Symbol::Function(FunctionSymbol {
            name: "len".to_string(),
            declaration: len_func,
            signature: len_signature,
            scope_level: 0,
            exported: true,
        });

        let _ = self.insert_global_symbol("len".to_string(), len_symbol);

        self.insert_builtin_assertions();
    }

    fn insert_builtin_assertions(&mut self) {
        use crate::ast::types::BasicType;
        use crate::error::Location;

        let base_location = Location::new(0, 0, 0);

        let builtin_asserts = vec![
            FunctionDecl {
                name: "assert".to_string(),
                generics: vec![],
                parameters: vec![
                    crate::ast::types::Parameter {
                        name: "condition".to_string(),
                        param_type: Type::Basic(BasicType::Bool),
                        default_value: None,
                        location: base_location,
                    },
                    crate::ast::types::Parameter {
                        name: "message".to_string(),
                        param_type: Type::Basic(BasicType::String),
                        default_value: None,
                        location: base_location,
                    },
                ],
                return_type: None,
                body: None,
                attributes: vec![],
                location: base_location,
            },
            FunctionDecl {
                name: "assert_true".to_string(),
                generics: vec![],
                parameters: vec![
                    crate::ast::types::Parameter {
                        name: "condition".to_string(),
                        param_type: Type::Basic(BasicType::Bool),
                        default_value: None,
                        location: base_location,
                    },
                    crate::ast::types::Parameter {
                        name: "message".to_string(),
                        param_type: Type::Basic(BasicType::String),
                        default_value: None,
                        location: base_location,
                    },
                ],
                return_type: None,
                body: None,
                attributes: vec![],
                location: base_location,
            },
            FunctionDecl {
                name: "assert_false".to_string(),
                generics: vec![],
                parameters: vec![
                    crate::ast::types::Parameter {
                        name: "condition".to_string(),
                        param_type: Type::Basic(BasicType::Bool),
                        default_value: None,
                        location: base_location,
                    },
                    crate::ast::types::Parameter {
                        name: "message".to_string(),
                        param_type: Type::Basic(BasicType::String),
                        default_value: None,
                        location: base_location,
                    },
                ],
                return_type: None,
                body: None,
                attributes: vec![],
                location: base_location,
            },
            FunctionDecl {
                name: "assert_eq".to_string(),
                generics: vec![],
                parameters: vec![
                    crate::ast::types::Parameter {
                        name: "actual".to_string(),
                        param_type: Type::Basic(BasicType::Any),
                        default_value: None,
                        location: base_location,
                    },
                    crate::ast::types::Parameter {
                        name: "expected".to_string(),
                        param_type: Type::Basic(BasicType::Any),
                        default_value: None,
                        location: base_location,
                    },
                    crate::ast::types::Parameter {
                        name: "message".to_string(),
                        param_type: Type::Basic(BasicType::String),
                        default_value: None,
                        location: base_location,
                    },
                ],
                return_type: None,
                body: None,
                attributes: vec![],
                location: base_location,
            },
            FunctionDecl {
                name: "assert_ne".to_string(),
                generics: vec![],
                parameters: vec![
                    crate::ast::types::Parameter {
                        name: "actual".to_string(),
                        param_type: Type::Basic(BasicType::Any),
                        default_value: None,
                        location: base_location,
                    },
                    crate::ast::types::Parameter {
                        name: "expected".to_string(),
                        param_type: Type::Basic(BasicType::Any),
                        default_value: None,
                        location: base_location,
                    },
                    crate::ast::types::Parameter {
                        name: "message".to_string(),
                        param_type: Type::Basic(BasicType::String),
                        default_value: None,
                        location: base_location,
                    },
                ],
                return_type: None,
                body: None,
                attributes: vec![],
                location: base_location,
            },
            FunctionDecl {
                name: "assert_panics".to_string(),
                generics: vec![],
                parameters: vec![
                    crate::ast::types::Parameter {
                        name: "test".to_string(),
                        param_type: Type::Function(crate::ast::types::FunctionType {
                            parameter_types: vec![],
                            return_type: None,
                            variadic: false,
                            location: base_location,
                        }),
                        default_value: None,
                        location: base_location,
                    },
                    crate::ast::types::Parameter {
                        name: "message".to_string(),
                        param_type: Type::Basic(BasicType::String),
                        default_value: None,
                        location: base_location,
                    },
                ],
                return_type: None,
                body: None,
                attributes: vec![],
                location: base_location,
            },
        ];

        for func in builtin_asserts {
            let signature = FunctionSignature::from_decl(&func);
            let symbol = Symbol::Function(FunctionSymbol {
                name: func.name.clone(),
                declaration: func.clone(),
                signature,
                scope_level: 0,
                exported: true,
            });
            let _ = self.insert_global_symbol(func.name.clone(), symbol);
        }
    }

    /// Enter a new scope
    pub fn enter_scope(&mut self) {
        self.current_level += 1;
        self.current_scope.push(HashMap::new());
    }

    /// Exit the current scope
    pub fn exit_scope(&mut self) {
        if self.current_level > 0 {
            self.current_scope.pop();
            self.current_level -= 1;
        }
    }

    /// Insert a function symbol
    pub fn insert_function(&mut self, func: &FunctionDecl) -> Result<()> {
        let signature = FunctionSignature::from_decl(func);
        let symbol = Symbol::Function(FunctionSymbol {
            name: func.name.clone(),
            declaration: func.clone(),
            signature,
            scope_level: self.current_level,
            exported: self.is_exported(&Declaration::Function(func.clone())),
        });

        self.insert_symbol(func.name.clone(), symbol)?;
        Ok(())
    }

    /// Insert a variable symbol
    pub fn insert_variable(&mut self, var: &VariableDecl) -> Result<()> {
        let symbol = Symbol::Variable(VariableSymbol {
            name: var.name.clone(),
            declaration: var.clone(),
            scope_level: self.current_level,
            exported: self.is_exported(&Declaration::Variable(var.clone())),
        });

        self.insert_symbol(var.name.clone(), symbol)?;
        Ok(())
    }

    /// Insert a constant symbol
    pub fn insert_constant(&mut self, const_: &ConstantDecl) -> Result<()> {
        let symbol = Symbol::Constant(ConstantSymbol {
            name: const_.name.clone(),
            declaration: const_.clone(),
            scope_level: self.current_level,
            exported: self.is_exported(&Declaration::Constant(const_.clone())),
        });

        self.insert_symbol(const_.name.clone(), symbol)?;
        Ok(())
    }

    /// Insert a type symbol
    pub fn insert_type(&mut self, type_: &TypeDecl) -> Result<()> {
        let symbol = Symbol::Type(TypeSymbol {
            name: type_.name.clone(),
            declaration: type_.clone(),
            scope_level: self.current_level,
            exported: self.is_exported(&Declaration::Type(type_.clone())),
        });

        self.insert_symbol(type_.name.clone(), symbol)?;
        Ok(())
    }

    /// Insert a struct symbol
    pub fn insert_struct(&mut self, struct_: &StructDecl) -> Result<()> {
        let symbol = Symbol::Struct(StructSymbol {
            name: struct_.name.clone(),
            declaration: struct_.clone(),
            scope_level: self.current_level,
            exported: self.is_exported(&Declaration::Struct(struct_.clone())),
        });

        self.insert_symbol(struct_.name.clone(), symbol)?;
        Ok(())
    }

    /// Insert an enum symbol
    pub fn insert_enum(&mut self, enum_: &EnumDecl) -> Result<()> {
        let symbol = Symbol::Enum(EnumSymbol {
            name: enum_.name.clone(),
            declaration: enum_.clone(),
            scope_level: self.current_level,
            exported: self.is_exported(&Declaration::Enum(enum_.clone())),
        });

        self.insert_symbol(enum_.name.clone(), symbol)?;
        Ok(())
    }

    /// Insert a trait symbol
    pub fn insert_trait(&mut self, trait_: &TraitDecl) -> Result<()> {
        let symbol = Symbol::Trait(TraitSymbol {
            name: trait_.name.clone(),
            declaration: trait_.clone(),
            scope_level: self.current_level,
            exported: self.is_exported(&Declaration::Trait(trait_.clone())),
        });

        self.insert_symbol(trait_.name.clone(), symbol)?;
        Ok(())
    }

    /// Insert an import symbol
    pub fn insert_import(&mut self, import: &ImportDecl) -> Result<()> {
        let symbol = Symbol::Import(ImportSymbol {
            path: import.path.clone(),
            declaration: import.clone(),
            scope_level: self.current_level,
            exported: self.is_exported(&Declaration::Import(import.clone())),
        });

        self.insert_symbol(import.path.clone(), symbol)?;
        Ok(())
    }

    /// Insert a symbol into the current scope
    fn insert_symbol(&mut self, name: String, symbol: Symbol) -> Result<()> {
        let scope = &mut self.current_scope[self.current_level];
        let entry = scope.entry(name.clone()).or_insert_with(Vec::new);
        Self::push_symbol_entry(entry, &name, symbol)
    }

    /// Insert a symbol into the global symbol map (used for builtins/imports)
    fn insert_global_symbol(&mut self, name: String, symbol: Symbol) -> Result<()> {
        Self::insert_symbol_into_map(&mut self.global_symbols, name, symbol)
    }

    /// Helper: insert symbol into arbitrary name -> Vec<Symbol> map
    fn insert_symbol_into_map(
        map: &mut HashMap<String, Vec<Symbol>>,
        name: String,
        symbol: Symbol,
    ) -> Result<()> {
        let entry = map.entry(name.clone()).or_insert_with(Vec::new);
        Self::push_symbol_entry(entry, &name, symbol)
    }

    /// Helper: push symbol into entry while enforcing overload rules
    fn push_symbol_entry(entry: &mut Vec<Symbol>, name: &str, symbol: Symbol) -> Result<()> {
        match &symbol {
            Symbol::Function(func_symbol) => {
                let duplicate = entry.iter().any(|existing| match existing {
                    Symbol::Function(existing_func) => {
                        existing_func.signature == func_symbol.signature
                    }
                    _ => false,
                });
                if duplicate {
                    return Err(CompilerError::semantic(
                        func_symbol.declaration.location.line,
                        func_symbol.declaration.location.column,
                        format!(
                            "Function '{}' with identical signature already declared in current scope",
                            name
                        ),
                    ));
                }
                entry.push(symbol);
            }
            _ => {
                if !entry.is_empty() {
                    return Err(CompilerError::semantic(
                        symbol.location().line,
                        symbol.location().column,
                        format!("Symbol '{}' already declared in current scope", name),
                    ));
                }
                entry.push(symbol);
            }
        }
        Ok(())
    }

    /// Look up a symbol by name
    pub fn lookup(&self, name: &str) -> Option<&Symbol> {
        // Search from current scope to global scope
        for scope in self.current_scope.iter().rev() {
            if let Some(symbols) = scope.get(name) {
                if let Some(symbol) = symbols.last() {
                    return Some(symbol);
                }
            }
        }

        // Check global symbols
        self.global_symbols
            .get(name)
            .and_then(|symbols| symbols.last())
    }

    /// Look up a symbol in the current scope only
    pub fn lookup_current_scope(&self, name: &str) -> Option<&Symbol> {
        self.current_scope[self.current_level]
            .get(name)
            .and_then(|symbols| symbols.last())
    }

    /// Look up the most recent function symbol with the given name
    pub fn lookup_function(&self, name: &str) -> Option<&FunctionSymbol> {
        self.lookup_function_in_scopes(name)
            .or_else(|| self.lookup_function_in_globals(name))
    }

    /// Get all overloads for a given function name (innermost scope first)
    pub fn lookup_function_overloads(&self, name: &str) -> Vec<&FunctionSymbol> {
        let mut results = Vec::new();
        for scope in self.current_scope.iter().rev() {
            if let Some(symbols) = scope.get(name) {
                for symbol in symbols.iter().rev() {
                    if let Symbol::Function(func) = symbol {
                        results.push(func);
                    }
                }
                if !results.is_empty() {
                    return results;
                }
            }
        }
        if let Some(symbols) = self.global_symbols.get(name) {
            for symbol in symbols.iter().rev() {
                if let Symbol::Function(func) = symbol {
                    results.push(func);
                }
            }
        }
        results
    }

    fn lookup_function_in_scopes(&self, name: &str) -> Option<&FunctionSymbol> {
        for scope in self.current_scope.iter().rev() {
            if let Some(symbols) = scope.get(name) {
                if let Some(symbol) = symbols.iter().rev().find_map(|symbol| {
                    if let Symbol::Function(func) = symbol {
                        Some(func)
                    } else {
                        None
                    }
                }) {
                    return Some(symbol);
                }
            }
        }
        None
    }

    fn lookup_function_in_globals(&self, name: &str) -> Option<&FunctionSymbol> {
        self.global_symbols.get(name).and_then(|symbols| {
            symbols.iter().rev().find_map(|symbol| {
                if let Symbol::Function(func) = symbol {
                    Some(func)
                } else {
                    None
                }
            })
        })
    }

    /// Look up a variable symbol
    pub fn lookup_variable(&self, name: &str) -> Option<&VariableSymbol> {
        if let Some(Symbol::Variable(var)) = self.lookup(name) {
            Some(var)
        } else {
            None
        }
    }

    /// Look up a constant symbol
    pub fn lookup_constant(&self, name: &str) -> Option<&ConstantSymbol> {
        if let Some(Symbol::Constant(const_)) = self.lookup(name) {
            Some(const_)
        } else {
            None
        }
    }

    /// Look up a type symbol
    pub fn lookup_type(&self, name: &str) -> Option<&TypeSymbol> {
        if let Some(Symbol::Type(type_)) = self.lookup(name) {
            Some(type_)
        } else {
            None
        }
    }

    /// Look up a struct symbol
    pub fn lookup_struct(&self, name: &str) -> Option<&StructSymbol> {
        if let Some(Symbol::Struct(struct_)) = self.lookup(name) {
            Some(struct_)
        } else {
            None
        }
    }

    /// Look up an enum symbol
    pub fn lookup_enum(&self, name: &str) -> Option<&EnumSymbol> {
        if let Some(Symbol::Enum(enum_)) = self.lookup(name) {
            Some(enum_)
        } else {
            None
        }
    }

    /// Look up a trait symbol
    pub fn lookup_trait(&self, name: &str) -> Option<&TraitSymbol> {
        if let Some(Symbol::Trait(trait_)) = self.lookup(name) {
            Some(trait_)
        } else {
            None
        }
    }

    /// Look up an import symbol
    pub fn lookup_import(&self, name: &str) -> Option<&ImportSymbol> {
        if let Some(Symbol::Import(import)) = self.lookup(name) {
            Some(import)
        } else {
            None
        }
    }

    /// Get all symbols in the current scope
    pub fn current_scope_symbols(&self) -> &HashMap<String, Vec<Symbol>> {
        &self.current_scope[self.current_level]
    }

    /// Get all global symbols
    pub fn global_symbols(&self) -> &HashMap<String, Vec<Symbol>> {
        &self.global_symbols
    }

    /// Get the current scope level
    pub fn current_level(&self) -> usize {
        self.current_level
    }

    /// Check if the symbol table is empty
    pub fn is_empty(&self) -> bool {
        let globals_empty = self
            .global_symbols
            .values()
            .all(|symbols| symbols.is_empty());
        let scopes_empty = self
            .current_scope
            .iter()
            .all(|scope| scope.values().all(|symbols| symbols.is_empty()));
        globals_empty && scopes_empty
    }

    /// Get the number of symbols in the current scope
    pub fn current_scope_size(&self) -> usize {
        self.current_scope[self.current_level]
            .values()
            .map(|symbols| symbols.len())
            .sum()
    }

    /// Get the total number of symbols
    pub fn total_symbols(&self) -> usize {
        let scope_symbols: usize = self
            .current_scope
            .iter()
            .map(|scope| scope.values().map(|symbols| symbols.len()).sum::<usize>())
            .sum();
        let global_symbols: usize = self
            .global_symbols
            .values()
            .map(|symbols| symbols.len())
            .sum();
        global_symbols + scope_symbols
    }

    /// Check if a declaration is exported
    fn is_exported(&self, decl: &Declaration) -> bool {
        match decl {
            Declaration::Function(func) => func.attributes.iter().any(|attr| attr.name == "export"),
            Declaration::Variable(_var) => {
                // Variables are not exported by default
                false
            }
            Declaration::Constant(_const_) => {
                // Constants are not exported by default
                false
            }
            Declaration::Type(_type_) => {
                // Types are not exported by default
                false
            }
            Declaration::Struct(_struct_) => {
                // Structs are not exported by default
                false
            }
            Declaration::Enum(_enum) => {
                // Enums are not exported by default
                false
            }
            Declaration::Trait(_trait) => {
                // Traits are not exported by default
                false
            }
            Declaration::TraitImpl(_trait_impl) => {
                // Trait implementations are not exported by default
                false
            }
            Declaration::Import(_import) => {
                // Imports are not exported by default
                false
            }
            Declaration::Impl(_impl) => {
                // Impl blocks are not exported by default
                false
            }
        }
    }

    /// Check if a function declaration is exported
    #[allow(dead_code)]
    fn is_exported_function(&self, func: &FunctionDecl) -> bool {
        func.attributes.iter().any(|attr| attr.name == "export")
    }

    /// Check if a variable declaration is exported
    #[allow(dead_code)]
    fn is_exported_variable(&self, _var: &VariableDecl) -> bool {
        // Variables are not exported by default
        false
    }

    /// Check if a constant declaration is exported
    #[allow(dead_code)]
    fn is_exported_constant(&self, _const_: &ConstantDecl) -> bool {
        // Constants are not exported by default
        false
    }

    /// Check if a type declaration is exported
    #[allow(dead_code)]
    fn is_exported_type(&self, _type_: &TypeDecl) -> bool {
        // Types are not exported by default
        false
    }

    /// Check if a struct declaration is exported
    #[allow(dead_code)]
    fn is_exported_struct(&self, _struct_: &StructDecl) -> bool {
        // Structs are not exported by default
        false
    }

    /// Check if an import declaration is exported
    #[allow(dead_code)]
    fn is_exported_import(&self, _import: &ImportDecl) -> bool {
        // Imports are not exported by default
        false
    }
}

impl Default for SymbolTable {
    fn default() -> Self {
        Self::new()
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_symbol_table_creation() {
        let table = SymbolTable::new();
        // Symbol table is not empty because it contains built-in functions
        assert_eq!(table.current_level(), 0);
    }

    #[test]
    fn test_scope_management() {
        let mut table = SymbolTable::new();

        assert_eq!(table.current_level(), 0);

        table.enter_scope();
        assert_eq!(table.current_level(), 1);

        table.exit_scope();
        assert_eq!(table.current_level(), 0);
    }

    #[test]
    fn test_symbol_insertion() {
        let mut table = SymbolTable::new();

        let func = FunctionDecl {
            name: "test".to_string(),
            generics: vec![],
            parameters: vec![],
            return_type: None,
            body: None,
            attributes: vec![],
            location: crate::error::Location::new(1, 1, 0),
        };

        table.insert_function(&func).unwrap();
        assert!(!table.is_empty());
        assert_eq!(table.current_scope_size(), 1);

        let symbol = table.lookup_function("test");
        assert!(symbol.is_some());
        assert_eq!(symbol.unwrap().name, "test");
    }

    #[test]
    fn test_symbol_lookup() {
        let mut table = SymbolTable::new();

        let func = FunctionDecl {
            name: "test".to_string(),
            generics: vec![],
            parameters: vec![],
            return_type: None,
            body: None,
            attributes: vec![],
            location: crate::error::Location::new(1, 1, 0),
        };

        table.insert_function(&func).unwrap();

        // Test function lookup
        let func_symbol = table.lookup_function("test");
        assert!(func_symbol.is_some());

        // Test general lookup
        let symbol = table.lookup("test");
        assert!(symbol.is_some());

        // Test non-existent symbol
        let non_existent = table.lookup("nonexistent");
        assert!(non_existent.is_none());
    }

    #[test]
    fn test_duplicate_symbol_error() {
        let mut table = SymbolTable::new();

        let func1 = FunctionDecl {
            name: "test".to_string(),
            generics: vec![],
            parameters: vec![],
            return_type: None,
            body: None,
            attributes: vec![],
            location: crate::error::Location::new(1, 1, 0),
        };

        let func2 = FunctionDecl {
            name: "test".to_string(),
            generics: vec![],
            parameters: vec![],
            return_type: None,
            body: None,
            attributes: vec![],
            location: crate::error::Location::new(2, 1, 0),
        };

        table.insert_function(&func1).unwrap();

        let result = table.insert_function(&func2);
        assert!(result.is_err());
    }
}
